import os

from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
import data_queries as data_queries

class SQLOutput(BaseModel):
    sql_query: str

model = os.getenv('PYDANTIC_AI_MODEL', 'gemini-1.5-flash')
print(f'Using model: {model}')


system_prompt = """
You are a SQL agent that can generate MSSQL queries based on natural language input. 
You should always use WITH(NOLOCK) in your queries.
Use tools to get table names and table ddls to help you generate the SQL queries.
"""

sql_agent = Agent(model, output_type=SQLOutput, system_prompt=system_prompt)

@sql_agent.tool_plain
def get_tables() -> list[str]:
    """
    Returns a list of table names in the database.
    """
    return data_queries.get_tables()

@sql_agent.tool
def get_table_schema(ctx: RunContext[str], table_name: str) -> str:
    """
    Returns the DDL schema of the specified table.
    """
    return data_queries.get_table_schema(table_name)

if __name__ == '__main__':
    prompt = """
    List all employees working in the 'Human Resources' department. 
    For each employee, display their full name (first name and last name), their job title, and the number of dependents they have. 
    Only include employees who have at least one dependent.
    The final output should contain the following columns: full_name, job_title, number_of_dependents.
    """
    result = sql_agent.run_sync(prompt)
    print(f"Generated SQL Query: {result.output.sql_query}")
